#!/usr/bin/env python3
import socket, json, threading, time, tkinter as tk
from tkinter import ttk, messagebox
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure
import numpy as np

SERVER = 'localhost'
PORT   = 5555
CANDLE_SECONDS = 60
MAX_CANDLES    = 100

# ---------------- 登录 ----------------
class LoginDialog(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("登录")
        self.geometry("280x120")
        self.resizable(False, False)
        self.username = None
        ttk.Label(self, text="用户名").pack(pady=5)
        self.ent = ttk.Entry(self, width=20)
        self.ent.pack()
        self.ent.focus()
        self.ent.bind("<Return>", lambda e: self.ok())
        ttk.Button(self, text="进入", command=self.ok).pack(pady=10)
    def ok(self):
        name = self.ent.get().strip()
        if name:
            self.username = name
            self.destroy()

# ---------------- 主窗口 ----------------
class MainWindow(tk.Tk):
    def __init__(self, username):
        super().__init__()
        self.title("杠杆炒股模拟器（独立仓位版）")
        self.geometry("1024x720")
        self.username = username
        self.sock = None
        self.price_history = []
        self.candles = []
        self.positions = []        # 当前仓位列表
        self.build_ui()
        self.connect()
        self.protocol("WM_DELETE_WINDOW", self.on_close)

    # ---------- 网络 ----------
    def connect(self):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((SERVER, PORT))
            self.sock.sendall((json.dumps({"type":"login","username":self.username})+"\n").encode())
            threading.Thread(target=self.recv_loop, daemon=True).start()
        except Exception as e:
            messagebox.showerror("连接失败", str(e))
            self.destroy()

    def recv_loop(self):
        reader = self.sock.makefile('r')
        while True:
            try:
                msg = json.loads(reader.readline())
                t = msg['type']
                if t == 'ticker':
                    price = msg['price']
                    ts = time.time()
                    self.price_history.append((ts, price))
                    self.price_lbl.config(text=f"价格: {price}")
                    self.online_lbl.config(text=f"在线: {msg['users']}")
                    self.update_candles()
                    # 更新仓位盈亏
                    if hasattr(self, 'state'):
                        self.update_assets()
                elif t == 'login_ok' or t == 'trade_done':
                    self.state = msg['state']
                    self.update_assets()
                elif t == 'chat':
                    self.chat_txt.config(state="normal")
                    self.chat_txt.insert("end", msg['msg'] + "\n")
                    self.chat_txt.config(state="disabled")
                    self.chat_txt.see("end")
                elif t == 'err':
                    messagebox.showerror("错误", msg['msg'])
            except Exception as e:
                messagebox.showerror("连接", f"与服务器断开\n{e}")
                self.destroy()
                break

    # ---------- UI ----------
    def build_ui(self):
        left  = ttk.Frame(self)
        right = ttk.Frame(self, width=280)
        left.pack(side="left", fill="both", expand=True)
        right.pack(side="right", fill="y")

        # 左：行情 + K线
        top = ttk.Frame(left)
        top.pack(fill="x", padx=5, pady=5)
        self.price_lbl = ttk.Label(top, text="价格: -", font=("Consolas", 16))
        self.price_lbl.pack(side="left")
        self.online_lbl = ttk.Label(top, text="在线: 0")
        self.online_lbl.pack(side="right")
        self.fig = Figure(figsize=(5,3), dpi=100)
        self.ax  = self.fig.add_subplot(111)
        self.canvas = FigureCanvasTkAgg(self.fig, master=left)
        self.canvas.get_tk_widget().pack(fill="both", expand=True, padx=5, pady=5)

        # 右：资产 + 开仓 + 转账 + 聊天 + 仓位
        ttk.Label(right, text="资产", font=("Consolas", 12)).pack(pady=3)
        self.asset_labels = {}
        for k in ("现金","未实现盈亏"):
            f = ttk.Frame(right)
            f.pack(fill="x", padx=5, pady=1)
            ttk.Label(f, text=k, width=10, anchor="w").pack(side="left")
            self.asset_labels[k] = ttk.Label(f, text="-", width=12)
            self.asset_labels[k].pack(side="left")

        # 开仓
        ttk.Separator(right, orient="horizontal").pack(fill="x", pady=5)
        ttk.Label(right, text="开仓").pack()
        frm = ttk.Frame(right); frm.pack()
        self.side_var = tk.StringVar(value="buy")
        ttk.Radiobutton(frm, text="买多", variable=self.side_var, value="buy").grid(row=0, column=0)
        ttk.Radiobutton(frm, text="卖空", variable=self.side_var, value="sell").grid(row=0, column=1)
        ttk.Label(frm, text="数量").grid(row=1, column=0)
        self.qty_ent = ttk.Entry(frm, width=8); self.qty_ent.grid(row=1, column=1)
        self.qty_ent.insert(0, "100")
        ttk.Label(frm, text="杠杆").grid(row=2, column=0)
        self.lev_scl = ttk.Scale(frm, from_=1, to=10, orient="horizontal", length=100)
        self.lev_scl.set(5); self.lev_scl.grid(row=2, column=1)
        self.lev_lbl = ttk.Label(frm, text="5", width=3)
        self.lev_lbl.grid(row=2, column=2)
        self.lev_scl.bind("<Motion>", lambda e: self.lev_lbl.config(text=str(int(self.lev_scl.get()))))
        ttk.Button(frm, text="开仓", command=self.send_open).grid(row=3, column=0, columnspan=3, pady=5)

        # 转账
        ttk.Separator(right, orient="horizontal").pack(fill="x", pady=5)
        ttk.Label(right, text="转账").pack()
        ttk.Label(right, text="给").pack()
        self.to_ent = ttk.Entry(right, width=12); self.to_ent.pack()
        ttk.Label(right, text="金额").pack()
        self.amt_ent = ttk.Entry(right, width=12); self.amt_ent.pack()
        ttk.Button(right, text="转账", command=self.send_transfer).pack(pady=3)

        # 聊天
        ttk.Separator(right, orient="horizontal").pack(fill="x", pady=5)
        ttk.Label(right, text="聊天室").pack()
        self.chat_txt = tk.Text(right, width=30, height=8, state="disabled")
        self.chat_txt.pack(fill="both", expand=True, padx=5, pady=5)
        self.msg_ent = ttk.Entry(right)
        self.msg_ent.pack(fill="x", padx=5)
        self.msg_ent.bind("<Return>", lambda e: self.send_chat())

        # 仓位列表
        ttk.Separator(right, orient="horizontal").pack(fill="x", pady=5)
        ttk.Label(right, text="当前仓位（双击平仓）").pack()
        self.pos_listbox = tk.Listbox(right, height=6)
        self.pos_listbox.pack(fill="both", expand=True, padx=5, pady=2)
        self.pos_listbox.bind("<Double-Button-1>", self.on_close_position)

    # ---------- K 线 ----------
    def update_candles(self):
        if not self.price_history:
            return
        ts0 = self.price_history[0][0]
        now = self.price_history[-1][0]
        candles = []
        start = ts0
        while start < now + CANDLE_SECONDS:
            end = start + CANDLE_SECONDS
            ticks = [p for t,p in self.price_history if start <= t < end]
            if ticks:
                o,h,l,c = ticks[0], max(ticks), min(ticks), ticks[-1]
                candles.append((start, o,h,l,c))
            start = end
        self.candles = candles[-MAX_CANDLES:]
        self.draw_candles()

    def draw_candles(self):
        self.ax.clear()
        if not self.candles:
            return
        times, opens, highs, lows, closes = zip(*self.candles)
        n = len(times)
        x = np.arange(n)
        width = max(0.05, 0.8 / max(n, 5))
        colors = ["green" if c >= o else "red" for o,c in zip(opens, closes)]
        for i,(o,h,l,c) in enumerate(zip(opens, highs, lows, closes)):
            self.ax.plot([x[i],x[i]],[l,h], color="black", linewidth=1)
            bottom = min(o,c); height = abs(c-o)
            self.ax.bar(x[i], height, bottom=bottom, width=width,
                        color=colors[i], edgecolor="black")
        self.ax.set_xlim(-0.5, max(n,5)-0.5)
        step = max(1, n//6)
        self.ax.set_xticks(x[::step])
        self.ax.set_xticklabels(
            [time.strftime("%H:%M", time.localtime(t)) for t in times[::step]],
            rotation=45, ha="right")
        self.ax.set_title(f"1 min K Line ({n})")
        self.ax.grid(True, linestyle="--", alpha=0.4)
        self.canvas.draw()

    # ---------- 发送 ----------
    def send_open(self):
        try:
            size = int(self.qty_ent.get())
            lev  = int(self.lev_scl.get())
            self.sock.sendall((json.dumps({
                "type":"trade_open",
                "side":self.side_var.get(),
                "size":size,
                "leverage":lev
            })+"\n").encode())
        except ValueError:
            messagebox.showerror("输入","数量/杠杆须为整数")

    def send_transfer(self):
        try:
            to = self.to_ent.get().strip()
            amt = int(self.amt_ent.get())
            self.sock.sendall((json.dumps({
                "type":"transfer",
                "to":to,
                "amount":amt
            })+"\n").encode())
        except ValueError:
            messagebox.showerror("输入","金额须为整数")

    def send_chat(self):
        msg = self.msg_ent.get().strip()
        if msg:
            self.sock.sendall((json.dumps({
                "type":"chat",
                "msg":msg
            })+"\n").encode())
            self.msg_ent.delete(0, "end")

    def on_close_position(self, event):
        sel = self.pos_listbox.curselection()
        if not sel:
            return
        pos_id = self.positions[sel[0]]['id']
        self.sock.sendall((json.dumps({
            "type":"trade_close",
            "id":pos_id
        })+"\n").encode())

    def update_assets(self):
        for k,lbl in self.asset_labels.items():
            lbl.config(text=str(int(self.state[k])))
        # 刷新仓位列表
        self.positions = self.state.get('仓位列表', [])
        self.pos_listbox.delete(0, tk.END)
        
        # 获取当前价格
        current_price = None
        if self.price_history:
            current_price = self.price_history[-1][1]
        
        for p in self.positions:
            side = "多" if p['size'] > 0 else "空"
            size = abs(p['size'])
            price = p['price']
            pos_id = p['id']
            
            # 计算盈亏
            pnl = 0
            if current_price is not None:
                if p['size'] > 0:  # 多头仓位
                    pnl = (current_price - price) * size
                else:  # 空头仓位
                    pnl = (price - current_price) * size
            
            # 显示仓位信息，包括盈亏
            pnl_str = f"盈亏: {pnl:+.2f}"
            self.pos_listbox.insert(tk.END,
                f"{side} {size} @ {price} {pnl_str} (id:{pos_id})")

    def on_close(self):
        if self.sock:
            self.sock.close()
        self.destroy()

# ---------------- 启动 ----------------
if __name__ == '__main__':
    login = LoginDialog()
    login.mainloop()
    if login.username:
        MainWindow(login.username).mainloop()
